/***************************************************************************
 *   Copyright (C) 2007 by Eduardo Robles Elvira   *
 *   edulix@gmail.com   *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include <errno.h>
#include <QTimer>
#include <QDebug>
#include <QMetaType>
#include <QDomDocument>
#include <iostream>
#include <QModelIndex>
#include <QItemSelectionModel>
#include <QVariant>
#include "konntrackmodel.h"
#include "konntrackitem.h"


using namespace std;


KonntrackModel::KonntrackModel(QObject *parent)
	: QAbstractTableModel(parent)
{
	qRegisterMetaType<QModelIndex>("QModelIndex");
	nextRow = 0;
	family = AF_INET;

	h = nfct_open(CONNTRACK, 0);
	if (!h) {
		qDebug() << "nfct_open error" << endl;
		perror("nfct_open");
		exit(-1);
	} 
	
	nfct_callback_register(h, NFCT_T_ALL, KonntrackModel::dump_cb, this);
	
	// timer gives a shot every second that updates the list of conntracks
	timer = new QTimer( this );
	connect( timer, SIGNAL(timeout()), this, SLOT(conntracksPopulate()) );
	timer->start( 1000 );
	conntracksPopulate(); // give it a first shot
}

 KonntrackModel::~KonntrackModel()
{
	nfct_close(h);
}

void KonntrackModel::setSelectionModel(QItemSelectionModel* selectionModel)
{
	this->selectionModel = selectionModel;
}

void  KonntrackModel::removeSelection()
{
// 	qDebug() << "void  KonntrackModel::removeSelection()" << endl;
	QModelIndexList list = selectionModel->selectedRows();
	QListIterator<QModelIndex> i(list);
	int row;
	while (i.hasNext()) {
		row = i.next().row();
		
		beginRemoveRows(QModelIndex(), row, row);
		
		KonntrackItem* item = items[row];
// 		qDebug() << "remove row: " << row << endl;
		/*int ret = */item->removeConntrack(h);
// 		qDebug() << "returns: " << ret << "error?" << strerror(errno) << endl;
		items.remove(row);
		delete item;
		
		endRemoveRows();
	}
}

void  KonntrackModel::ipNameResolved()
{
	emit dataChanged(index(0,0), index(items.size()-1,19));
}


void KonntrackModel::conntracksPopulate()
{
	nextRow = 0;
	nfct_query(h, NFCT_Q_DUMP, &family);
	
	//Set num rows
	if(items.size() > nextRow) {
		beginRemoveRows(QModelIndex(), nextRow+1, items.size());
		for(int i = nextRow; i  < items.size(); i++)
		{
			delete items[i];
		}
		items.resize(nextRow);
		endRemoveRows();
	}
	
	emit dataChanged(index(0,0), index(nextRow,19));
}

void KonntrackModel::listConntrack(char string[], nf_conntrack *ct)
{
	bool resized = false;
	if(items.size() <= nextRow)
	{
		beginInsertRows(QModelIndex(), nextRow, nextRow);
		items.resize(nextRow+1);
		items[nextRow] = new KonntrackItem;
		connect(items[nextRow], SIGNAL(ipNameResolved()), this, SLOT(ipNameResolved()));
		resized = true;
	}
	
	items[nextRow]->setConntrack(ct);
	
	QString buffer(string);
	QDomDocument domTree;
	if ( !domTree.setContent( buffer ) ) {
		qDebug() << "domTree error" << endl;
	}
	QDomElement root = domTree.documentElement();
	
	QDomNode node = root.firstChild();
	while ( !node.isNull() ) {
		if ( node.isElement() && node.nodeName() == "meta" ) {
			QDomElement meta = node.toElement();
			processMeta( meta );
		}
		node = node.nextSibling();
	}
	
	if(resized)
		endInsertRows();
	
	nextRow++;
}


void KonntrackModel::setColText(int col, QDomNode &node)
{
	QDomText textChild = node.firstChild().toText();
	if ( !textChild.isNull() ) {
		items[nextRow]->setData((KonntrackCol)col, textChild.nodeValue());
	}
}

void KonntrackModel::processMeta(QDomElement &meta)
{
	int type=0; // type 0 means direction, 1 means independent
	int col=0;
	if(meta.attribute("direction") == "original") {
		type = 0;
	} else if(meta.attribute("direction") == "reply") {
		type = 0;
		col = 7; // col starts in reply columns
	} else if(meta.attribute("direction") == "independent") {
		type = 1;
	} else {
		qDebug() << "unknown attribute: " << meta.attribute("direction") << endl;
		return;
	}
	if(type == 0) {
		QDomNode node = meta.firstChild();
		QString proto[2] = {"", ""};
		while ( !node.isNull() ) {
			if ( node.isElement() ) {
				if ( node.nodeName() == "layer3" ) {
					proto[0] = node.toElement().attribute( "protoname" );
					QDomNode node2 = node.firstChild();
					while ( !node2.isNull() ) {
						if ( node2.isElement() ) {
							if( node2.nodeName() == "src" )
								setColText(col + 1, node2);
							else if( node2.nodeName() == "dst" )
								setColText(col + 3, node2);
						}
						node2 = node2.nextSibling();
					}
				} else if ( node.nodeName() == "layer4" ) {
					proto[1] = node.toElement().attribute( "protoname" );
					QDomNode node2 = node.firstChild();
					while ( !node2.isNull() ) {
						if ( node2.isElement() ) {
							if( node2.nodeName() == "sport" )
								setColText(col + 2, node2);
							else if( node2.nodeName() == "dport" )
								setColText(col + 4, node2);
						}
						node2 = node2.nextSibling();
					}
				} else if ( node.nodeName() == "counters" ) {
					QDomNode node2 = node.firstChild();
					while ( !node2.isNull() ) {
						if ( node2.isElement() ) {
							if( node2.nodeName() == "packets" )
								setColText(col + 5, node2);
							else if( node2.nodeName() == "bytes" )
								setColText(col + 6, node2);
						}
						node2 = node2.nextSibling();
					}
				}
			}
			node = node.nextSibling();
		}
		items[nextRow]->setData((KonntrackCol)col, proto[0] + "/" + proto[1]);
	} else if(type == 1) {
		QDomNode node = meta.firstChild();
		bool assured = false, unreplied = false;
		while ( !node.isNull() ) {
			if ( node.isElement() ) {
				if( node.nodeName() == "timeout" )
					setColText(14, node);
				else if( node.nodeName() == "mark" )
					setColText(15, node);
				else if( node.nodeName() == "use" )
					setColText(16, node);
				else if( node.nodeName() == "assured" ) {
					items[nextRow]->setData((KonntrackCol)17, "true");
					assured = true;
				}
				else if( node.nodeName() == "unreplied" ) {
					items[nextRow]->setData((KonntrackCol)18, "true");
					assured = true;
				}
			}
			node = node.nextSibling();
		}
		
		if(!assured)
			items[nextRow]->setData((KonntrackCol)17, "false");
		
		if(!unreplied)
			items[nextRow]->setData((KonntrackCol)18, "false");
	}
}




int KonntrackModel::dump_cb(enum nf_conntrack_msg_type /*type*/, struct nf_conntrack *ct, void *data)
{
	KonntrackModel *self = (KonntrackModel *)data;
	char buf[1024];

	nfct_snprintf(buf, 1024, ct, NFCT_T_UNKNOWN, NFCT_O_XML, 0);
	
	self->listConntrack(buf, ct);
	
	return NFCT_CB_STOLEN;
}

QVariant KonntrackModel::data(const QModelIndex &index, int role) const
{
	return data(index, role, true);
}

 QVariant KonntrackModel::data(const QModelIndex &index, int role, bool translated) const
{	
	if (!index.isValid())
		return QVariant();

	if (role != Qt::DisplayRole)
		return QVariant();

	return items[index.row()]->getData(index.column(), translated);
}

QVariant KonntrackModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	if (orientation == Qt::Horizontal && role == Qt::DisplayRole) {
		switch (section) {
			case ORIG_PROTO:
				return tr("Orig protocols");
			case ORIG_SRC_IP:
				return tr("Orig src");
			case ORIG_SPORT:
				return tr("Orig sport");
			case ORIG_DST_IP:
				return tr("Orig dst ip");
			case ORIG_DPORT:
				return tr("Orig dport");
			case ORIG_PACKETS:
				return tr("Orig packets");
			case ORIG_BYTES:
				return tr("Orig bytes");
			case REPLY_PROTO:
				return tr("Orig protocols");
			case REPLY_SRC_IP:
				return tr("Orig src ip");
			case REPLY_SPORT:
				return tr("Orig sport");
			case REPLY_DST_IP:
				return tr("Orig dst ip");
			case REPLY_DPORT:
				return tr("Orig dport");
			case REPLY_PACKETS:
				return tr("Orig packets");
			case REPLY_BYTES:
				return tr("Orig bytes");
			case TIMEOUT:
				return tr("Timeout");
			case MARK:
				return tr("Mark");
			case USE:
				return tr("Use");
			case ASSURED:
				return tr("Assured");
			case UNREPLIED:
				return tr("Unreplied");
			case NUMBER:
				return tr("No.");
			default:
				return QVariant();
		}
	}
	// never reach here
	return QVariant();
}

int KonntrackModel::rowCount(const QModelIndex &/*parent*/) const
{
	return items.size();
}

int KonntrackModel::columnCount(const QModelIndex &/*parent*/) const
{
	return 20;
}



